#include "nets.hpp"

CustomDtaset::CustomDtaset(at::Tensor input_data, at::Tensor label_data)
{
    inputs_ = input_data;
    labels_ = label_data;
}

torch::data::Example<> CustomDtaset::get(size_t index)
{
    torch::Tensor Input = inputs_[index];
    torch::Tensor Label = labels_[index];
    return {Input.clone(), Label.clone()};
}

torch::optional<size_t> CustomDtaset::size() const { return labels_.size(0); }

Net::Net(const int inDim, const int outDim)
{
    const int hiddenSize = 1;
    phase = register_module("phase", torch::nn::Linear(inDim, hiddenSize));
    trans = register_module("trans", torch::nn::Linear(hiddenSize, outDim));
}

at::Tensor Net::forward(at::Tensor t)
{
    t = phase->forward(t);
    t = torch::cos(t);
    t = trans->forward(t);
    return t;
}
